import math
import numpy as np
import python.version as version
from python.helpers import (
    add_vec,
    euclidean_dist,
    mat_mul,
    reduce_sym_list,
    sub_vec,
    norm,
    most_likely_list,
    print_v,
)

# Inefficient but enough for testing
from python.simple_gauss import simple_gauss

BP_alone = []
new_lwe_instance = []
dim = []
distance_between = []
Usvp_Basis = []
Lattice_volum = []
smallest_vector = []
range_of_BKZ = []
cofficients = []
result = []


def substitute_s(lwe_instance, known_s):
    lwe = lwe_instance.copy()
    for j, s_j in known_s:
        for i, (a_i, b_i) in enumerate(zip(lwe.a, lwe.b)):
            lwe.b[i] = (b_i - a_i[j] * s_j) % lwe.q
            a_i[j] = 0
    return lwe


def change_is_solution(a, b, new_key, mapping_e, q):
    m = len(a[0])  # 列数
    n = version.KYBER_N
    assert all(
        (len(ai) == m for ai in a)
    ), "Matrix 'a' does not have consistent row lengths"
    s = new_key[-version.KYBER_k * version.KYBER_N :]  # 负数最多到-1
    res = []
    for index, rows in enumerate(mapping_e[: -version.KYBER_k * version.KYBER_N]):
        exchange = [
            sum(items)
            for items in zip(*[a[rows + k * n] for k in range(version.KYBER_k)])
        ]
        result_left = sum(e * f for e, f in zip(exchange, s)) % q
        result_right = (
            sum(b[rows + k * n] for k in range(version.KYBER_k)) - new_key[index]
        ) % q
        res.append(result_left == result_right)
    return all(res)


def changes_is_solution(a, b, new_key, mapping_e, q, count_e):  # 进行removed_s之后
    n = version.KYBER_N
    m = len(a[0])  # 列数
    assert all(
        (len(ai) == m for ai in a)
    ), "Matrix 'a' does not have consistent row lengths"
    e_sum = new_key[:count_e]
    s = new_key[count_e:]  # 负数最多到-1
    print_v(f"count_e == len(mapping_e)-512? {count_e == len(mapping_e) - 512}")
    print_v(f"列数k == 长度s ? {m == len(s)}")
    print_v(m, len(s))
    res = []
    count = 0
    for index, rows in enumerate(mapping_e[: -version.KYBER_k * version.KYBER_N]):
        exchange = [
            sum(items)
            for items in zip(*[a[rows + k * n] for k in range(version.KYBER_k)])
        ]
        result_left = sum(e * f for e, f in zip(exchange, s)) % q
        result_right = (
            sum(b[rows + k * n] for k in range(version.KYBER_k)) - new_key[index]
        ) % q
        res.append(result_left == result_right)
    return all(res)


def substitute_e(lwe_instance, known_e, key=None):
    # 我们其实不需要这个eliminate_order_s，因为对e的消去不同，可以不以s作为函数输入
    q = lwe_instance.q
    lwe = lwe_instance.copy()
    n = version.KYBER_N
    m = len(lwe.a[0])
    substitution_matrix = []
    substitution_matrix_vec = []
    rows_to_delete = []
    rows_to_known_e = []

    for i, e_i in known_e:
        substitution_matrix_indices = [i + n * k for k in range(1, version.KYBER_k)]
        substitution_matrix.append(
            lwe.a[idx].copy() for idx in substitution_matrix_indices
        )
        substitution_matrix_vec.append(
            (lwe.b[idx]) % q for idx in substitution_matrix_indices
        )
        for index in substitution_matrix_indices:
            lwe.a[i] = [a + b for a, b in zip(lwe.a[i], lwe.a[index])]
            lwe.b[i] = lwe.b[i] + lwe.b[index]
            lwe.a[index] = [0] * m
            lwe.b[index] = 0
        lwe.b[i] = (lwe.b[i] - e_i) % q
        rows_to_known_e.append(i)
        rows_to_delete.extend(substitution_matrix_indices)

        if key:
            result_left = (
                sum(x * y for x, y in zip(lwe.a[i], key[version.KYBER_N :])) % q
            )
            assert result_left == lwe.b[i]
    new_key = []
    mapping_e = []
    for index, value in enumerate(key):
        if (
            index not in rows_to_known_e
        ):  # 对于known_e中的结果我们实际都测试过了，是能通过assert的
            new_key.append(value)
            mapping_e.append(index)
    assert change_is_solution(lwe.a, lwe.b, new_key, mapping_e, q)

    return (
        lwe,
        (substitution_matrix, substitution_matrix_vec),
        rows_to_delete,  # 恢复出来的e的个数，但实际都是后256行里的
        mapping_e,
        new_key,
    )


def remove_solved_s(
    lwe_instance, del_list, mapping_e, key=None
):  # 这里是已经去掉一些e的key
    q = lwe_instance.q
    lwe = lwe_instance.copy()
    del_list = sorted(del_list)
    n = len(lwe.a)  # 行数
    m = len(lwe.a[0])  # 列数
    if key:
        assert change_is_solution(lwe.a, lwe.b, key, mapping_e, q)
    if not del_list:  # 当恢复的s的数量为空
        mapping_s = [None for _ in range(m)]
        new_key = key[: -version.KYBER_k * version.KYBER_N]
        count_e = len(new_key)
        return lwe, mapping_s, key, mapping_e, count_e
    current_to_remove = del_list.pop()  # 像栈一样一个个蹦出来
    delete_s = []
    mapping_s = []
    number = 0
    for j in reversed(range(m)):  # m是列数
        if j == current_to_remove:
            for i in range(n):
                assert lwe.a[i][j] == 0
                del lwe.a[i][j]
            if key:
                delete_s.append(j)  # 要删除的key中的s，删除第几个s
                number += 1
            if del_list:
                current_to_remove = del_list.pop()
    print_v(f" number = {number}")
    new_key = key[: -version.KYBER_k * version.KYBER_N]
    count_e = len(new_key)
    for index, value in enumerate(key[-version.KYBER_k * version.KYBER_N :]):
        if (
            index not in delete_s
        ):  # 对于known_e中的结果我们实际都测试过了，是能通过assert的
            new_key.append(value)
            mapping_s.append(
                index
            )  # mapping_s中存储的是key的s对应的真实索引（因为已经删除了部分）
    # 所以需要对mapping_e进行截断
    assert changes_is_solution(lwe.a, lwe.b, new_key, mapping_e, q, count_e)
    # 在还没有删除所有的0行向量时，再看一下是否满足,new_key已经是去掉恢复出来的e和s后的了
    return lwe, mapping_s, new_key, mapping_e, count_e
    # 还是没有对行进行约简，仅仅是设成了0，new_key是剩下的所有的未被恢复的e和s，对应的映射为mapping_e,mapping_s


def solve_from_substituted(known_s, substituted_equations, lwe_instance, key=None):
    q = lwe_instance.q
    n = len(substituted_equations[0][0])
    for j, s_j in known_s:
        substituted_equations[0].append([1 if k == j else 0 for k in range(n)])
        substituted_equations[1].append(s_j)
    ###
    if key:
        b = list(mat_mul(substituted_equations[0], key[512:], q))
        assert reduce_sym_list(b) == reduce_sym_list(substituted_equations[1])
    ###
    key_s = simple_gauss(substituted_equations[0], substituted_equations[1], q)
    key_s = list(reduce_sym_list(key_s, q))
    b = list(mat_mul(substituted_equations[0], key_s, q))
    assert reduce_sym_list(b) == reduce_sym_list(substituted_equations[1])
    if key:
        assert key[n:] == key_s
    a_s = list(mat_mul(lwe_instance.a, key_s, q))
    key_e = list(reduce_sym_list(sub_vec(lwe_instance.b, a_s, q)))
    if key:
        assert lwe_instance.is_solution(key)
        a_s_test = list(mat_mul(lwe_instance.a, key[n:], q))
        assert list(a_s) == list(a_s_test)
        e_test = list(reduce_sym_list(sub_vec(lwe_instance.b, a_s_test, q)))
        assert e_test == key_e
        as_e = list(reduce_sym_list(add_vec(a_s, key_e, q), q))
        assert as_e == list(reduce_sym_list(lwe_instance.b, q))
    rec_key = key_e + key_s
    if key:
        assert key == rec_key
    return rec_key


def solve(
    propagation_data,
    block_size,
    max_beta,
    perform,
    add_fplll,
    step,
    max_enum,
    step_rank,
):
    print_v("Solving LWE instance..")
    key = propagation_data.key
    n = version.KYBER_N
    if step == -1:
        current_step = sorted(propagation_data.steps.items(), key=lambda x: x[0])[step][
            1
        ]
    else:
        current_step = propagation_data.steps[step]
    recovered_positions = current_step.recovered_coefficients
    ordered_key_indices = current_step.ordered_key_indices
    known_indices = ordered_key_indices[:recovered_positions]
    known_e = []
    known_s = []
    for i in known_indices:
        if i < n:
            known_e.append((i, current_step.guessed_key[i]))
        else:
            known_s.append((i - n, current_step.guessed_key[i]))
    print_v(f"known_e: {known_e}\n")
    print_v(f"known_s: {known_s}")
    print_v(f"len(known_s)+len(known_e) is equal to {len(known_s) + len(known_e)}")
    assert all((propagation_data.key[i] == known_e_i for i, known_e_i in known_e))
    assert all((propagation_data.key[i + n] == known_s_i for i, known_s_i in known_s))

    print_v("Substituting known coefficients of s..")  # 函数终止在这里
    lwe_instance = substitute_s(
        propagation_data.lwe_instance, known_s
    )  # 对于恢复的knows，已经置为0
    assert lwe_instance.is_solution(key)
    print_v("Eliminating coefficients using known coefficients of e..")
    (
        lwe_instance,
        substituted_equations,
        removed_e,  # 就是恢复的e_sum的个数
        mapping_e,
        key_rem,
    ) = substitute_e(
        lwe_instance,
        known_e,
        propagation_data.key.copy(),
    )

    print_v("今天天气不错")
    solved_s_indices = list(map(lambda x: x[0], known_s))  # 恢复出来的密钥s的索引
    print_v(f"Recovered s with {len(solved_s_indices)} indices.")
    if (
        len(solved_s_indices) < version.KYBER_k * version.KYBER_N
    ):  # 原来是n，这里有待商榷
        zero_row_count = 0
        for row in lwe_instance.a:
            if all(element == 0 for element in row):
                zero_row_count += 1
        print_v(
            f"Preparing new lwe instance with dimension {(len(lwe_instance.a) - zero_row_count, len(lwe_instance.a[0]))}.."
        )

        print_v("Removing solved s..")
        lwe, mapping_s, key_rem, mapping_e, count_e = remove_solved_s(
            lwe_instance, solved_s_indices, mapping_e, key_rem
        )  # 这个要改一下

        print_v(
            f"New lwe instance has dim(a) = {(len(lwe.a) - zero_row_count, len(lwe.a[0]))}."
        )  #

        new_lwe_instance.append(len(lwe_instance.a))
        new_lwe_instance.append(len(lwe_instance.a[0]) + 1)
        result.append(new_lwe_instance)
        # assert lwe.is_solution(key_rem)
        # assert len(removed_e) == len(set(removed_e))
        # print_v(f"New lwe instance has dim(a) = {(len(lwe.a), len(lwe.a[0]))}.")

        dim.append(len(lwe.a) - zero_row_count)
        dim.append(len(lwe.a[0]))
        result.append(dim)

        remaining_indices = [i for i in range(n) if i + n not in removed_e] + [
            i + n
            for i in range(version.KYBER_k * version.KYBER_N)
            if i not in solved_s_indices
        ]  # 不知道的e的数量+不知道的s的数量，removed_e里放的是我们删去的（置为零的行的索引，是后256行）
        print_v(f"remaining_indices is equal to {remaining_indices}")
        print_v(f"mapping_e[:count_e] = {mapping_e[:count_e]}")
        print_v(f"mapping_s = {mapping_s}")
        print_v(768 - 328 == len(remaining_indices))
        print_v(256 - len(known_e) + 512 - len(known_s), len(remaining_indices))
        print_v(f"len(known_e) == count_e? {len(known_e) == count_e}")
        dist_key_remaining = [
            dist
            for i, dist in enumerate(current_step.results)
            if i in remaining_indices
        ]
        print_v(f"dist_key_remaining is equal to {dist_key_remaining}")
        (
            success,
            res,
            usvp_basis,
            bikz,
            guess_rank,
            key_rank,
            distance,
        ) = estimate_hinted_lwe(
            lwe,
            propagation_data,
            dist_key_remaining,
            key_rem,
            max_beta,
            perform,
            block_size,
            add_fplll,
            max_enum,
            mapping_e,
            count_e,
        )
        propagation_data.set_lattice_data(
            usvp_basis,
            bikz,
            guess_rank,
            step,
            step_rank,
            key_rank,
            distance,
        )
        if not success:
            return False
        if not perform:
            return True

        for idx, coeff in zip(remaining_indices, res):
            assert key[idx] == coeff or key_rem[idx] == -coeff
            if idx >= n:
                solved_s_indices.append(idx)
                known_s.append((idx - n, coeff))

        assert len(solved_s_indices) == n or not success
    else:
        propagation_data.set_lattice_data(None, (0, 0), 0, step, step_rank, 0, 0)
        if not perform:
            return True

    if len(solved_s_indices) >= n:
        for i, si in known_s:
            assert key[i + n] == si
        if perform:
            key = solve_from_substituted(
                known_s,
                substituted_equations,
                propagation_data.lwe_instance,
                propagation_data.key,
            )
            assert key == propagation_data.key
        return True
    return False


# More general function for tests, not needed here
def ln_lattice_volume(basis):
    basis = np.array(basis)
    sq = basis.transpose() @ basis
    sign, log_det = np.linalg.slogdet(sq)
    return log_det / 2


def delta_beta(beta):
    delta = pow(math.pi * beta, 1 / beta)
    delta *= beta / (2 * math.pi * math.e)
    delta = pow(delta, 1 / (2 * beta - 2))
    return delta


def is_solveable(beta, dim, log_volume, norm_s):
    left = math.sqrt(beta / dim) * norm_s
    delta = delta_beta(beta)
    right = pow(delta, 2 * beta - dim - 1) * math.exp(log_volume / dim)
    return left <= right


# Inefficient but works
def get_bikz(dim, ln_volume, short, step=1, max_b=5000):
    norm_s = norm(short)
    for beta in range(50, max_b, step):
        if is_solveable(beta, dim, ln_volume, norm_s):
            if beta > 50:
                return beta, beta - step
            else:
                return beta, 0
    return None, max_b


def usvp_basis_from_lwe(lwe, guess=None):
    n = len(lwe.a)  # 行数
    m = len(lwe.a[0])  # 列数
    guess = guess if guess else [0 for _ in range(n + m)]
    print_v(f"guess is equal to {guess}")
    new_b = [bi for bi in lwe.b] + [0] * m + [1]
    print_v(f"new_b is equal to {new_b}")
    print_v(f"guess[: len(lwe.b)] is equal to {guess[: len(lwe.b)]}")
    print_v(f"guess[len(lwe.b):] is equal to {guess[len(lwe.b):]}")
    usvp_basis = [[0 for _ in range(n + m + 1)] for _ in range(n + m)]
    for i in range(n):
        usvp_basis[i][i] = lwe.q
        for j in range(m):
            usvp_basis[j + n][i] = lwe.a[i][j]
    for i in range(m):
        usvp_basis[i + n][i + n] = -1
    usvp_basis.append(new_b)
    return usvp_basis


def sanitize_dists(dists):
    dists_new = []
    for d in dists:
        dn = {}
        for v, p in d.items():
            if p > 0.01:
                dn[v] = p
        if len(dn) == 0:
            dn = d.copy()
        dists_new.append(dn)
    return dists_new


def get_best_key(dists, max_enum, key, num_bins=50, num_merged=20):
    from histo_guesser import HistogramGuesser

    print_v("Enumerating..")
    dists = sanitize_dists(dists)
    assert len(dists) == len(key)
    guesser = HistogramGuesser.build(dists, num_bins, num_merged)
    rk = guesser.key_estimation(key)
    if rk:
        print_v(f"Approximate key rank: 2^{math.log2(rk)}")
    else:
        print_v("Approximate key rank greater than 2^128")
    best = guesser.next_key()
    assert best is not None
    best_dist = euclidean_dist(best, key)
    best_guess_idx = 0
    print_v(f"Dist of first guess (sanity check): {best_dist}")
    for i in range(max_enum):
        guess = guesser.next_key()
        if guess is None:
            break
        dist = euclidean_dist(key, guess)
        if dist < best_dist:
            best = guess
            best_dist = dist
            best_guess_idx = i
    print_v(f"Best key at position {best_guess_idx}.")
    return best, best_guess_idx, rk


def estimate_hinted_lwe(
    lwe,
    propagation_data,
    dist_key_remaining,
    key_rem,
    max_beta,
    perform,
    block_size,
    add_fplll,
    max_enum,
    mapping_e,
    count_e,
):
    guess = list(most_likely_list(dist_key_remaining))
    dist = euclidean_dist(guess, key_rem)
    print_v(f"Distance between naive hint and key: {dist}")
    if max_enum > 1:
        guess, guess_rank, key_rank = get_best_key(
            dist_key_remaining, max_enum, key_rem
        )
        dist = euclidean_dist(guess, key_rem)
        print_v(f"Distance between best hint and key: {dist}")

    distance_between.append(dist)
    result.append(distance_between)

    key_rem = list(key_rem)
    guess_rank, key_rank = 0, None
    if dist == 0:
        print_v("Solved.")
    print_v(len(guess), len(key_rem))
    assert len(guess) == len(key_rem)  # 未恢复的e和s
    q = lwe.q
    assert changes_is_solution(lwe.a, lwe.b, key_rem, mapping_e, q, count_e)
    print_v("=================")
    rows_to_remove = []
    for i, row in enumerate(lwe.a):
        if all(elem == 0 for elem in row) and lwe.b[i] == 0:
            rows_to_remove.append(i)

        # 反向遍历需要删除的行索引，并从后往前删除，以避免索引问题
    for i in reversed(rows_to_remove):
        del lwe.a[i]
        del lwe.b[i]
    n = len(lwe.a)  # 行数
    m = len(lwe.a[0])  # 列数
    # with open('lwe_instance.txt', 'a') as file:  # 多余的行和列都已经被消去了
    #     file.write(f"lwe.a = {lwe.a}\n")
    #     file.write(f"lwe.b = {lwe.b}\n")
    usvp_basis = usvp_basis_from_lwe(lwe, guess)

    Usvp_Basis.append(len(usvp_basis))
    Usvp_Basis.append(len(usvp_basis[0]))
    result.append(Usvp_Basis)

    print_v(f"uSVP basis has dimensions {(len(usvp_basis), len(usvp_basis[0]))}")
    lat_elem = list(sub_vec(key_rem, guess))
    dim = len(usvp_basis)

    assert dim == n + m + 1
    assert all((len(b) == dim for b in usvp_basis))
    volume = n * math.log(lwe.q)
    print_v(f"Lattice volume: {volume}")
    Lattice_volum.append(volume)
    result.append(Lattice_volum)
    print_v(f"Smallest vector norm: {norm(lat_elem)}")

    smallest_vector.append(norm(lat_elem))
    result.append(smallest_vector)

    bikz_upper, bikz_lower = get_bikz(dim, volume, lat_elem)
    propagation_data.lattice_data = usvp_basis.copy()  # TODO: Deep copy
    propagation_data.bikz = (bikz_lower, bikz_upper)
    print_v(f"{bikz_lower} <= BIKZ <= {bikz_upper}")

    range_of_BKZ.append(bikz_lower)
    range_of_BKZ.append(bikz_upper)
    result.append(range_of_BKZ)

    # with open('output1.txt', 'a') as file:
    #     file.write(f"参数为：{result}\n")

    if bikz_upper <= max_beta or bikz_upper <= 50:
        if perform:
            res = bkz_fplll(
                usvp_basis,
                block_size if block_size else bikz_upper,
                propagation_data.get_dir(),
                add_fplll,
            )
            recovered_partial_key = list(add_vec(res, guess))
            return (
                True,
                recovered_partial_key,
                usvp_basis,
                (bikz_lower, bikz_upper),
                guess_rank,
                key_rank,
                dist,
            )
        return (
            True,
            None,
            usvp_basis,
            (bikz_lower, bikz_upper),
            guess_rank,
            key_rank,
            dist,
        )
    return (
        False,
        None,
        usvp_basis,
        (bikz_lower, bikz_upper),
        guess_rank,
        key_rank,
        dist,
    )


def bkz_fplll(basis, block_size, path, add_fplll):
    import subprocess
    import os

    # Strange hack because fpylll is difficult to install on some systems
    basis = str(basis).replace(",", "")
    os.makedirs("data", exist_ok=True)
    filename = f"{path}/fplll_mat.txt"
    with open(filename, "w") as f:
        f.write(basis)
    cmd = f"fplll {add_fplll} -a bkz -b {block_size} {filename}"
    # cmd = f"fplll -v -a lll {filename}"
    print_v(f"Running '{cmd}'")
    output = subprocess.check_output(cmd, shell=True)
    output = output.decode("ascii").replace("\n", ",").replace(" ", ", ")
    mat = eval(output)[0]
    return mat[0][:-1]
